import math
import torch
import torch.nn as nn


class PositionalEncoding(nn.Module):
    """

    Args:
        dropout : dropout
        dim : dim
        max_len : max_lengths
    Input:
        embedding : [batch, dim]
        step : None
    Output:
        embedding : [batch, dim]
    """
    def __init__(self, dropout, dim, max_len=5000):
        if dim % 2 != 0:
            raise ValueError("Cannot use sin/cos positional encoding with "
                             "odd dim (got dim={:d})".format(dim))
        pe = torch.zeros(max_len, dim)  # [max_len, dim]
        position = torch.arange(0, max_len).unsqueeze(1)  # [max_len, 1]
        div_term = torch.exp((torch.arange(0, dim, 2, dtype=torch.float) *
                              -(math.log(10000.0) / dim)))
        pe[:, 0::2] = torch.sin(position.float() * div_term)
        pe[:, 1::2] = torch.cos(position.float() * div_term)
        super(PositionalEncoding, self).__init__()
        self.register_buffer('pe', pe)
        self.dropout = nn.Dropout(p=dropout)
        self.dim = dim

    def forward(self, emb, step=None):
        emb = emb * math.sqrt(self.dim)
        step = step or 0
        if self.pe.size(0) < step + emb.size(1):
            raise ValueError(
                f"Sequence is {emb.size(1) + step} but PositionalEncoding is"
                f" limited to {self.pe.size(0)}. See max_len argument."
            )
        emb = emb + self.pe[step:emb.size(1) + step]
        emb = self.dropout(emb)
        return emb


if __name__ == '__main__':
    t_emb = torch.zeros(4, 10, 2)
    t_pe = PositionalEncoding(0.1, dim=2, max_len=10)
    t_out = t_pe(t_emb)
    print(t_out - t_emb)
